import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, GATConv, global_mean_pool
from config import *
from graph_utils import build_grid_graph


class GraphEncoder(nn.Module):
    """使用图神经网络进行状态编码的模块"""

    def __init__(self, node_features, edge_features, hidden_channels, num_layers):
        super(GraphEncoder, self).__init__()
        self.node_features = node_features
        self.edge_features = edge_features
        self.hidden_channels = hidden_channels
        self.num_layers = num_layers

        # 节点特征嵌入
        self.node_encoder = nn.Linear(node_features, hidden_channels)

        # 图卷积层
        self.conv_layers = nn.ModuleList()
        self.batch_norms = nn.ModuleList()

        # 第一层使用边特征
        self.conv_layers.append(GATConv(hidden_channels, hidden_channels))
        self.batch_norms.append(nn.BatchNorm1d(hidden_channels))

        # 添加更多卷积层
        for _ in range(num_layers - 1):
            self.conv_layers.append(GATConv(hidden_channels, hidden_channels))
            self.batch_norms.append(nn.BatchNorm1d(hidden_channels))

        # 创建固定的边索引和属性
        self.register_buffer('edge_index', None)
        self.register_buffer('edge_attr', None)

        # 初始化边索引和属性
        edge_index, edge_attr = build_grid_graph()
        self.register_buffer('edge_index', edge_index)
        self.register_buffer('edge_attr', edge_attr)

    def forward(self, x, batch=None):
        # 节点特征嵌入
        x = self.node_encoder(x)

        # 应用图卷积层
        for i in range(self.num_layers):
            x = self.conv_layers[i](x, self.edge_index)
            x = self.batch_norms[i](x)
            x = F.relu(x)

        # 如果提供了batch，则执行全局池化
        if batch is not None:
            x = global_mean_pool(x, batch)

        return x


class EnhancedWorldModel(nn.Module):
    """使用图神经网络和增强状态表示的改进世界模型"""

    def __init__(self, node_features, dropout_rate=DROPOUT_RATE):
        super(EnhancedWorldModel, self).__init__()

        # 特征维度
        self.node_features = node_features

        # 图编码器
        self.graph_encoder = GraphEncoder(
            node_features=node_features,
            edge_features=GNN_EDGE_FEATURES,
            hidden_channels=GNN_HIDDEN_CHANNELS,
            num_layers=GNN_NUM_LAYERS
        )

        # 全连接层 - 从图编码到环境和智能体预测
        grid_size = ENV_SIZE
        total_nodes = grid_size * grid_size
        encoded_dim = GNN_HIDDEN_CHANNELS * total_nodes

        # 共享层
        self.fc1 = nn.Linear(encoded_dim, HIDDEN_SIZES[0])
        self.bn1 = nn.BatchNorm1d(HIDDEN_SIZES[0])
        self.dropout1 = nn.Dropout(dropout_rate)

        # 环境状态预测分支
        self.env_fc1 = nn.Linear(HIDDEN_SIZES[0], HIDDEN_SIZES[1])
        self.env_bn1 = nn.BatchNorm1d(HIDDEN_SIZES[1])
        self.env_dropout1 = nn.Dropout(dropout_rate)

        self.env_fc2 = nn.Linear(HIDDEN_SIZES[1], HIDDEN_SIZES[2])
        self.env_bn2 = nn.BatchNorm1d(HIDDEN_SIZES[2])
        self.env_dropout2 = nn.Dropout(dropout_rate)

        self.env_output = nn.Linear(HIDDEN_SIZES[2], OUTPUT_SIZE)

        # 智能体预测分支 - 增加容量和深度
        self.agent_fc1 = nn.Linear(HIDDEN_SIZES[0], AGENT_BRANCH_HIDDEN_SIZES[0])
        self.agent_bn1 = nn.BatchNorm1d(AGENT_BRANCH_HIDDEN_SIZES[0])
        self.agent_dropout1 = nn.Dropout(dropout_rate)

        self.agent_fc2 = nn.Linear(AGENT_BRANCH_HIDDEN_SIZES[0], AGENT_BRANCH_HIDDEN_SIZES[1])
        self.agent_bn2 = nn.BatchNorm1d(AGENT_BRANCH_HIDDEN_SIZES[1])
        self.agent_dropout2 = nn.Dropout(dropout_rate)

        self.agent_fc3 = nn.Linear(AGENT_BRANCH_HIDDEN_SIZES[1], AGENT_BRANCH_HIDDEN_SIZES[2])
        self.agent_bn3 = nn.BatchNorm1d(AGENT_BRANCH_HIDDEN_SIZES[2])
        self.agent_dropout3 = nn.Dropout(dropout_rate)

        self.agent_fc4 = nn.Linear(AGENT_BRANCH_HIDDEN_SIZES[2], AGENT_BRANCH_HIDDEN_SIZES[3])
        self.agent_bn4 = nn.BatchNorm1d(AGENT_BRANCH_HIDDEN_SIZES[3])
        self.agent_dropout4 = nn.Dropout(dropout_rate)

        # 添加残差连接
        self.agent_residual = nn.Linear(AGENT_BRANCH_HIDDEN_SIZES[0], AGENT_BRANCH_HIDDEN_SIZES[3])

        self.agent_fc5 = nn.Linear(AGENT_BRANCH_HIDDEN_SIZES[3], AGENT_BRANCH_HIDDEN_SIZES[4])
        self.agent_bn5 = nn.BatchNorm1d(AGENT_BRANCH_HIDDEN_SIZES[4])
        self.agent_dropout5 = nn.Dropout(dropout_rate)

        # 输出层 - 预测相对移动(dx,dy)和绝对位置(x,y)
        self.agent_output = nn.Linear(AGENT_BRANCH_HIDDEN_SIZES[4], 4)

    def reshape_for_graph(self, x):
        """将批次输入重塑为图形式"""
        batch_size = x.size(0)
        grid_size = ENV_SIZE

        # 假设输入x的形状是[batch_size, grid_size*grid_size*CELL_FEATURES]
        # 重塑为[batch_size, grid_size, grid_size, CELL_FEATURES]
        x = x.view(batch_size, grid_size, grid_size, -1)

        # 现在将其转换为适合图的形式
        # 对于每个批次样本，创建图节点特征
        node_features = []
        for b in range(batch_size):
            # 对于每个单元格，提取特征
            sample_nodes = []
            for i in range(grid_size):
                for j in range(grid_size):
                    # 获取单元格状态特征
                    cell_features = x[b, i, j].tolist()

                    # 添加位置编码
                    pos_i = 2 * (i / (grid_size - 1)) - 1  # 范围从-1到1
                    pos_j = 2 * (j / (grid_size - 1)) - 1  # 范围从-1到1

                    # 添加位置编码到特征中
                    node_feature = cell_features + [pos_i, pos_j]
                    sample_nodes.append(node_feature)

            node_features.append(sample_nodes)

        # 转换为适当的形状: [batch_size * grid_size * grid_size, node_features]
        node_features = torch.tensor(node_features, dtype=torch.float, device=x.device)
        node_features = node_features.view(batch_size * grid_size * grid_size, -1)

        return node_features

    def forward(self, x):
        """前向传递"""
        batch_size = x.size(0)
        grid_size = ENV_SIZE
        total_nodes = grid_size * grid_size

        # 将输入重塑为图形式
        node_features = self.reshape_for_graph(x)

        # 创建batch索引，用于区分不同样本的节点
        batch_idx = []
        for b in range(batch_size):
            batch_idx.extend([b] * total_nodes)
        batch_idx = torch.tensor(batch_idx, dtype=torch.long, device=x.device)

        # 使用图编码器处理节点特征
        node_encodings = self.graph_encoder(node_features)

        # 将节点编码重塑为批次形式
        node_encodings = node_encodings.view(batch_size, total_nodes * GNN_HIDDEN_CHANNELS)

        # 共享层
        x = self.fc1(node_encodings)
        x = self.bn1(x)
        x = F.relu(x)
        shared_features = self.dropout1(x)

        # 环境状态预测分支
        env_x = self.env_fc1(shared_features)
        env_x = self.env_bn1(env_x)
        env_x = F.relu(env_x)
        env_x = self.env_dropout1(env_x)

        env_x = self.env_fc2(env_x)
        env_x = self.env_bn2(env_x)
        env_x = F.relu(env_x)
        env_x = self.env_dropout2(env_x)

        env_output = self.env_output(env_x)
        env_output = env_output.view(batch_size, ENV_SIZE, ENV_SIZE, CELL_FEATURES)

        # 智能体预测分支 - 使用更深层网络和残差连接
        agent_x = self.agent_fc1(shared_features)
        agent_x = self.agent_bn1(agent_x)
        agent_x = F.relu(agent_x)
        agent_x1 = self.agent_dropout1(agent_x)  # 保存用于残差连接

        agent_x = self.agent_fc2(agent_x1)
        agent_x = self.agent_bn2(agent_x)
        agent_x = F.relu(agent_x)
        agent_x = self.agent_dropout2(agent_x)

        agent_x = self.agent_fc3(agent_x)
        agent_x = self.agent_bn3(agent_x)
        agent_x = F.relu(agent_x)
        agent_x = self.agent_dropout3(agent_x)

        agent_x = self.agent_fc4(agent_x)
        agent_x = self.agent_bn4(agent_x)
        agent_x = F.relu(agent_x)
        agent_x = self.agent_dropout4(agent_x)

        # 添加残差连接
        agent_residual = self.agent_residual(agent_x1)
        agent_x = agent_x + agent_residual

        agent_x = self.agent_fc5(agent_x)
        agent_x = self.agent_bn5(agent_x)
        agent_x = F.relu(agent_x)
        agent_x = self.agent_dropout5(agent_x)

        # 输出相对移动和绝对位置
        agent_output = self.agent_output(agent_x)  # (dx, dy, x, y)

        return env_output, agent_output


# 保留原始世界模型类，以便向后兼容
class WorldModel(nn.Module):
    """具有双重输出的世界模型"""

    def __init__(self, dropout_rate=DROPOUT_RATE):
        super(WorldModel, self).__init__()
        # 输入层
        self.fc1 = nn.Linear(INPUT_SIZE, HIDDEN_SIZES[0])
        self.bn1 = nn.BatchNorm1d(HIDDEN_SIZES[0])
        self.dropout1 = nn.Dropout(dropout_rate)

        # 中间层
        self.fc2 = nn.Linear(HIDDEN_SIZES[0], HIDDEN_SIZES[1])
        self.bn2 = nn.BatchNorm1d(HIDDEN_SIZES[1])
        self.dropout2 = nn.Dropout(dropout_rate)

        self.fc3 = nn.Linear(HIDDEN_SIZES[1], HIDDEN_SIZES[2])
        self.bn3 = nn.BatchNorm1d(HIDDEN_SIZES[2])
        self.dropout3 = nn.Dropout(dropout_rate)

        # 分支1：预测环境状态
        self.env_branch = nn.Linear(HIDDEN_SIZES[2], OUTPUT_SIZE)

        # 分支2：额外关注智能体位置预测
        self.agent_branch = nn.Linear(HIDDEN_SIZES[2], 4)  # (dx, dy, x, y)

    def forward(self, x):
        # 确保输入正确的批次维度
        if x.dim() == 1:
            x = x.unsqueeze(0)

        batch_size = x.size(0)

        # 共享编码器部分
        x = self.fc1(x)
        x = x.view(batch_size, -1)
        x = self.bn1(x)
        x = F.relu(x)
        x = self.dropout1(x)

        x = self.fc2(x)
        x = self.bn2(x)
        x = F.relu(x)
        x = self.dropout2(x)

        x = self.fc3(x)
        x = self.bn3(x)
        x = F.relu(x)
        x = self.dropout3(x)

        # 环境状态预测
        env_output = self.env_branch(x)
        env_output = env_output.view(batch_size, ENV_SIZE, ENV_SIZE, CELL_FEATURES)

        # 智能体位置预测
        agent_output = self.agent_branch(x)

        return env_output, agent_output